import json
from collections.abc import Iterable, Iterator

import numpy as np
from frozendict import frozendict
from zkl_pyutils_fsspec import Fs, FsLike, resolve_fs
from zkl_pyutils_serialization import dump_json_value, parse_json_value

from .homo_bytes import HomoBytesSequenceReader, HomoBytesSequenceWriter
from .loader import RecordJson, record_json_file_name
from .mergeable import Mergeable


class NumpyRecordReader(Mergeable):
    def __init__(self, fs: Fs, *, shape: Iterable[int], dtype: np.dtype | str):
        self._shape = tuple(map(int, shape))
        self._dtype = np.dtype(dtype)
        item_size = self._dtype.itemsize * int(np.prod(self._shape))
        self._reader = HomoBytesSequenceReader(fs, item_size=item_size)

    def _convert(self, item: bytes) -> np.ndarray:
        item = np.frombuffer(item, dtype=self._dtype)
        item = np.reshape(item, self._shape)
        return item

    @property
    def merging_signature(self) -> frozendict:
        return frozendict(type='ndarray', shape=self._shape, dtype=self._dtype.str)

    def refresh(self, clean: bool):
        self._reader.refresh(clean)

    def __len__(self) -> int:
        return len(self._reader)

    def __iter__(self) -> Iterator[np.ndarray]:
        return map(self._convert, self._reader)

    def __getitem__(self, item_i: int) -> np.ndarray:
        item = self._reader[item_i]
        return self._convert(item)


class NumpyRecordWriter:
    def __init__(self, fs: FsLike, *, shape: Iterable[int], dtype: np.dtype | str):
        fs = resolve_fs(fs)
        fs.makedirs("", exist_ok=True)

        self._fs = fs
        self._shape = tuple(shape)
        self._dtype = np.dtype(dtype)
        item_size = self._dtype.itemsize * int(np.prod(self._shape))
        self._writer = HomoBytesSequenceWriter(self._fs, item_size=item_size)

        if self._fs.exists(record_json_file_name):
            with self._fs.open(record_json_file_name, 'rt', encoding='utf-8') as fp:
                record_json = parse_json_value(json.load(fp), RecordJson)
            read_shape = record_json.kwargs['shape']
            read_shape = tuple(map(int, read_shape))
            read_dtype = record_json.kwargs.get('dtype')
            read_dtype = np.dtype(read_dtype)
            if read_shape != self._shape:
                raise ValueError(f'shape mismatch: expected {read_shape}, got {self._shape}')
            if read_dtype != self._dtype:
                raise ValueError(f'dtype mismatch: expected {read_dtype}, got {self._dtype}')
        else:
            record_json = RecordJson.of(NumpyRecordReader, shape=self._shape, dtype=self._dtype.str)
            with self._fs.open(record_json_file_name, 'wt', encoding='utf-8') as fp:
                json.dump(dump_json_value(record_json), fp, indent=2, ensure_ascii=False)

    def _convert(self, item: np.ndarray) -> bytes:
        item = np.asarray(item, dtype=self._dtype)
        if item.shape != self._shape:
            raise ValueError(f'shape mismatch: expected {self._shape}, got {item.shape}')
        return item.tobytes()

    def write(self, item: np.ndarray):
        item = self._convert(item)
        self._writer.write(item)
